import torch 
from DenoriseDiffusion import DoubleUnet, Diffusion
from matplotlib import pyplot as plt
%matplotlib inline
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def im_show(arr, row, col):
    fig, axs = plt.subplots(nrows=row, ncols=col)
    for r in range(row):
        for c in range(col):
            axs[r][c].imshow(arr[r*col+c])
            axs[r][c].axis('off')
    plt.show()
diffusion = torch.load("./model/DoubleDiffusion_mnist.pth", map_location=device).to(device)
diffusion.device = device
images = diffusion.sample(16)
imgs = images.cpu().numpy().transpose(0,2,3,1)
im_show(imgs, 4,4)